{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# How to add semantic search to your agent's memory\n",
    "\n",
    "This guide shows how to enable semantic search in your agent's memory store. This lets search for items in the store by semantic similarity.\n",
    "\n",
    "!!! tip Prerequisites\n",
    "    This guide assumes familiarity with the [memory in LangGraph](https://langchain-ai.github.io/langgraph/concepts/memory/).\n",
    "\n",
    "First, install this guide's prerequisites."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture --no-stderr\n",
    "%pip install -U langgraph langchain-openai langchain"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import getpass\n",
    "import os\n",
    "\n",
    "\n",
    "def _set_env(var: str):\n",
    "    if not os.environ.get(var):\n",
    "        os.environ[var] = getpass.getpass(f\"{var}: \")\n",
    "\n",
    "\n",
    "_set_env(\"OPENAI_API_KEY\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, create the store with an [index configuration](https://langchain-ai.github.io/langgraph/reference/store/#langgraph.store.base.IndexConfig). By default, stores are configured without semantic/vector search. You can opt in to indexing items when creating the store by providing an [IndexConfig](https://langchain-ai.github.io/langgraph/reference/store/#langgraph.store.base.IndexConfig) to the store's constructor. If your store class does not implement this interface, or if you do not pass in an index configuration, semantic search is disabled, and all `index` arguments passed to `put` or `aput` will have no effect. Below is an example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/var/folders/gf/6rnp_mbx5914kx7qmmh7xzmw0000gn/T/ipykernel_83572/2318027494.py:5: LangChainBetaWarning: The function `init_embeddings` is in beta. It is actively being worked on, so the API may change.\n",
      "  embeddings = init_embeddings(\"openai:text-embedding-3-small\")\n"
     ]
    }
   ],
   "source": [
    "from langchain.embeddings import init_embeddings\n",
    "from langgraph.store.memory import InMemoryStore\n",
    "\n",
    "# Create store with semantic search enabled\n",
    "embeddings = init_embeddings(\"openai:text-embedding-3-small\")\n",
    "store = InMemoryStore(\n",
    "    index={\n",
    "        \"embed\": embeddings,\n",
    "        \"dims\": 1536,\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's store some memories:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Store some memories\n",
    "store.put((\"user_123\", \"memories\"), \"1\", {\"text\": \"I love pizza\"})\n",
    "store.put((\"user_123\", \"memories\"), \"2\", {\"text\": \"I prefer Italian food\"})\n",
    "store.put((\"user_123\", \"memories\"), \"3\", {\"text\": \"I don't like spicy food\"})\n",
    "store.put((\"user_123\", \"memories\"), \"3\", {\"text\": \"I am studying econometrics\"})\n",
    "store.put((\"user_123\", \"memories\"), \"3\", {\"text\": \"I am a plumber\"})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Search memories using natural language:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Memory: I prefer Italian food (similarity: 0.46482669521168163)\n",
      "Memory: I love pizza (similarity: 0.35514845174380766)\n",
      "Memory: I am a plumber (similarity: 0.155698702336571)\n"
     ]
    }
   ],
   "source": [
    "# Find memories about food preferences\n",
    "memories = store.search((\"user_123\", \"memories\"), query=\"I like food?\", limit=5)\n",
    "\n",
    "for memory in memories:\n",
    "    print(f\"Memory: {memory.value['text']} (similarity: {memory.score})\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using in your agent\n",
    "\n",
    "Add semantic search to any node by injecting the store."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "What are you in the mood for? Since you love Italian food and pizza, would you like to order a pizza or try making one at home?"
     ]
    }
   ],
   "source": [
    "from typing import Optional\n",
    "\n",
    "from langchain.chat_models import init_chat_model\n",
    "from langgraph.store.base import BaseStore\n",
    "\n",
    "from langgraph.graph import START, MessagesState, StateGraph\n",
    "\n",
    "llm = init_chat_model(\"openai:gpt-4o-mini\")\n",
    "\n",
    "\n",
    "def chat(state, *, store: BaseStore):\n",
    "    # Search based on user's last message\n",
    "    items = store.search(\n",
    "        (\"user_123\", \"memories\"), query=state[\"messages\"][-1].content, limit=2\n",
    "    )\n",
    "    memories = \"\\n\".join(item.value[\"text\"] for item in items)\n",
    "    memories = f\"## Memories of user\\n{memories}\" if memories else \"\"\n",
    "    response = llm.invoke(\n",
    "        [\n",
    "            {\"role\": \"system\", \"content\": f\"You are a helpful assistant.\\n{memories}\"},\n",
    "            *state[\"messages\"],\n",
    "        ]\n",
    "    )\n",
    "    return {\"messages\": [response]}\n",
    "\n",
    "\n",
    "builder = StateGraph(MessagesState)\n",
    "builder.add_node(chat)\n",
    "builder.add_edge(START, \"chat\")\n",
    "graph = builder.compile(store=store)\n",
    "\n",
    "for message, metadata in graph.stream(\n",
    "    input={\"messages\": [{\"role\": \"user\", \"content\": \"I'm hungry\"}]},\n",
    "    stream_mode=\"messages\",\n",
    "):\n",
    "    print(message.content, end=\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using in `create_react_agent` {#using-in-create-react-agent}\n",
    "\n",
    "Add semantic search to your tool calling agent by injecting the store in the `prompt` function. You can also use the store in a tool to let your agent manually store or search for memories."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "import uuid\n",
    "from typing import Optional\n",
    "\n",
    "from langchain.chat_models import init_chat_model\n",
    "from langgraph.prebuilt import InjectedStore\n",
    "from langgraph.store.base import BaseStore\n",
    "from typing_extensions import Annotated\n",
    "\n",
    "from langgraph.prebuilt import create_react_agent\n",
    "\n",
    "\n",
    "def prepare_messages(state, *, store: BaseStore):\n",
    "    # Search based on user's last message\n",
    "    items = store.search(\n",
    "        (\"user_123\", \"memories\"), query=state[\"messages\"][-1].content, limit=2\n",
    "    )\n",
    "    memories = \"\\n\".join(item.value[\"text\"] for item in items)\n",
    "    memories = f\"## Memories of user\\n{memories}\" if memories else \"\"\n",
    "    return [\n",
    "        {\"role\": \"system\", \"content\": f\"You are a helpful assistant.\\n{memories}\"}\n",
    "    ] + state[\"messages\"]\n",
    "\n",
    "\n",
    "# You can also use the store directly within a tool!\n",
    "def upsert_memory(\n",
    "    content: str,\n",
    "    *,\n",
    "    memory_id: Optional[uuid.UUID] = None,\n",
    "    store: Annotated[BaseStore, InjectedStore],\n",
    "):\n",
    "    \"\"\"Upsert a memory in the database.\"\"\"\n",
    "    # The LLM can use this tool to store a new memory\n",
    "    mem_id = memory_id or uuid.uuid4()\n",
    "    store.put(\n",
    "        (\"user_123\", \"memories\"),\n",
    "        key=str(mem_id),\n",
    "        value={\"text\": content},\n",
    "    )\n",
    "    return f\"Stored memory {mem_id}\"\n",
    "\n",
    "\n",
    "agent = create_react_agent(\n",
    "    init_chat_model(\"openai:gpt-4o-mini\"),\n",
    "    tools=[upsert_memory],\n",
    "    # The 'prompt' function is run to prepare the messages for the LLM. It is called\n",
    "    # right before each LLM call\n",
    "    prompt=prepare_messages,\n",
    "    store=store,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "What are you in the mood for? Since you love Italian food and pizza, maybe something in that realm would be great! Would you like suggestions for a specific dish or restaurant?"
     ]
    }
   ],
   "source": [
    "for message, metadata in agent.stream(\n",
    "    input={\"messages\": [{\"role\": \"user\", \"content\": \"I'm hungry\"}]},\n",
    "    stream_mode=\"messages\",\n",
    "):\n",
    "    print(message.content, end=\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Advanced Usage\n",
    "\n",
    "#### Multi-vector indexing\n",
    "\n",
    "Store and search different aspects of memories separately to improve recall or omit certain fields from being indexed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Expect mem 2\n",
      "Item: mem2; Score (0.5895009051396596)\n",
      "Memory: Ate alone at home\n",
      "Emotion: felt a bit lonely\n",
      "\n",
      "Expect mem1\n",
      "Item: mem1; Score (0.6207546534134083)\n",
      "Memory: Had pizza with friends at Mario's\n",
      "Emotion: felt happy and connected\n",
      "\n",
      "Expect random lower score (ravioli not indexed)\n",
      "Item: mem1; Score (0.2686278787315685)\n",
      "Memory: Had pizza with friends at Mario's\n",
      "Emotion: felt happy and connected\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Configure store to embed both memory content and emotional context\n",
    "store = InMemoryStore(\n",
    "    index={\"embed\": embeddings, \"dims\": 1536, \"fields\": [\"memory\", \"emotional_context\"]}\n",
    ")\n",
    "# Store memories with different content/emotion pairs\n",
    "store.put(\n",
    "    (\"user_123\", \"memories\"),\n",
    "    \"mem1\",\n",
    "    {\n",
    "        \"memory\": \"Had pizza with friends at Mario's\",\n",
    "        \"emotional_context\": \"felt happy and connected\",\n",
    "        \"this_isnt_indexed\": \"I prefer ravioli though\",\n",
    "    },\n",
    ")\n",
    "store.put(\n",
    "    (\"user_123\", \"memories\"),\n",
    "    \"mem2\",\n",
    "    {\n",
    "        \"memory\": \"Ate alone at home\",\n",
    "        \"emotional_context\": \"felt a bit lonely\",\n",
    "        \"this_isnt_indexed\": \"I like pie\",\n",
    "    },\n",
    ")\n",
    "\n",
    "# Search focusing on emotional state - matches mem2\n",
    "results = store.search(\n",
    "    (\"user_123\", \"memories\"), query=\"times they felt isolated\", limit=1\n",
    ")\n",
    "print(\"Expect mem 2\")\n",
    "for r in results:\n",
    "    print(f\"Item: {r.key}; Score ({r.score})\")\n",
    "    print(f\"Memory: {r.value['memory']}\")\n",
    "    print(f\"Emotion: {r.value['emotional_context']}\\n\")\n",
    "\n",
    "# Search focusing on social eating - matches mem1\n",
    "print(\"Expect mem1\")\n",
    "results = store.search((\"user_123\", \"memories\"), query=\"fun pizza\", limit=1)\n",
    "for r in results:\n",
    "    print(f\"Item: {r.key}; Score ({r.score})\")\n",
    "    print(f\"Memory: {r.value['memory']}\")\n",
    "    print(f\"Emotion: {r.value['emotional_context']}\\n\")\n",
    "\n",
    "print(\"Expect random lower score (ravioli not indexed)\")\n",
    "results = store.search((\"user_123\", \"memories\"), query=\"ravioli\", limit=1)\n",
    "for r in results:\n",
    "    print(f\"Item: {r.key}; Score ({r.score})\")\n",
    "    print(f\"Memory: {r.value['memory']}\")\n",
    "    print(f\"Emotion: {r.value['emotional_context']}\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Override fields at storage time\n",
    "You can override which fields to embed when storing a specific memory using `put(..., index=[...fields])`, regardless of the store's default configuration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Expect mem1\n",
      "Item: mem1; Score (0.3374968677940555)\n",
      "Memory: I love spicy food\n",
      "Context: At a Thai restaurant\n",
      "\n",
      "Expect mem2\n",
      "Item: mem2; Score (0.36784461593247436)\n",
      "Memory: The restaurant was too loud\n",
      "Context: Dinner at an Italian place\n",
      "\n"
     ]
    }
   ],
   "source": [
    "store = InMemoryStore(\n",
    "    index={\n",
    "        \"embed\": embeddings,\n",
    "        \"dims\": 1536,\n",
    "        \"fields\": [\"memory\"],\n",
    "    }  # Default to embed memory field\n",
    ")\n",
    "\n",
    "# Store one memory with default indexing\n",
    "store.put(\n",
    "    (\"user_123\", \"memories\"),\n",
    "    \"mem1\",\n",
    "    {\"memory\": \"I love spicy food\", \"context\": \"At a Thai restaurant\"},\n",
    ")\n",
    "\n",
    "# Store another overriding which fields to embed\n",
    "store.put(\n",
    "    (\"user_123\", \"memories\"),\n",
    "    \"mem2\",\n",
    "    {\"memory\": \"The restaurant was too loud\", \"context\": \"Dinner at an Italian place\"},\n",
    "    index=[\"context\"],  # Override: only embed the context\n",
    ")\n",
    "\n",
    "# Search about food - matches mem1 (using default field)\n",
    "print(\"Expect mem1\")\n",
    "results = store.search(\n",
    "    (\"user_123\", \"memories\"), query=\"what food do they like\", limit=1\n",
    ")\n",
    "for r in results:\n",
    "    print(f\"Item: {r.key}; Score ({r.score})\")\n",
    "    print(f\"Memory: {r.value['memory']}\")\n",
    "    print(f\"Context: {r.value['context']}\\n\")\n",
    "\n",
    "# Search about restaurant atmosphere - matches mem2 (using overridden field)\n",
    "print(\"Expect mem2\")\n",
    "results = store.search(\n",
    "    (\"user_123\", \"memories\"), query=\"restaurant environment\", limit=1\n",
    ")\n",
    "for r in results:\n",
    "    print(f\"Item: {r.key}; Score ({r.score})\")\n",
    "    print(f\"Memory: {r.value['memory']}\")\n",
    "    print(f\"Context: {r.value['context']}\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Disable Indexing for Specific Memories\n",
    "\n",
    "Some memories shouldn't be searchable by content. You can disable indexing for these while still storing them using \n",
    "`put(..., index=False)`. Example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Expect mem1\n",
      "Item: mem1; Score (0.32269984224327286)\n",
      "Memory: I love chocolate ice cream\n",
      "Type: preference\n",
      "\n",
      "Expect low score (mem2 not indexed)\n",
      "Item: mem1; Score (0.010241633698527089)\n",
      "Memory: I love chocolate ice cream\n",
      "Type: preference\n",
      "\n"
     ]
    }
   ],
   "source": [
    "store = InMemoryStore(index={\"embed\": embeddings, \"dims\": 1536, \"fields\": [\"memory\"]})\n",
    "\n",
    "# Store a normal indexed memory\n",
    "store.put(\n",
    "    (\"user_123\", \"memories\"),\n",
    "    \"mem1\",\n",
    "    {\"memory\": \"I love chocolate ice cream\", \"type\": \"preference\"},\n",
    ")\n",
    "\n",
    "# Store a system memory without indexing\n",
    "store.put(\n",
    "    (\"user_123\", \"memories\"),\n",
    "    \"mem2\",\n",
    "    {\"memory\": \"User completed onboarding\", \"type\": \"system\"},\n",
    "    index=False,  # Disable indexing entirely\n",
    ")\n",
    "\n",
    "# Search about food preferences - finds mem1\n",
    "print(\"Expect mem1\")\n",
    "results = store.search((\"user_123\", \"memories\"), query=\"what food preferences\", limit=1)\n",
    "for r in results:\n",
    "    print(f\"Item: {r.key}; Score ({r.score})\")\n",
    "    print(f\"Memory: {r.value['memory']}\")\n",
    "    print(f\"Type: {r.value['type']}\\n\")\n",
    "\n",
    "# Search about onboarding - won't find mem2 (not indexed)\n",
    "print(\"Expect low score (mem2 not indexed)\")\n",
    "results = store.search((\"user_123\", \"memories\"), query=\"onboarding status\", limit=1)\n",
    "for r in results:\n",
    "    print(f\"Item: {r.key}; Score ({r.score})\")\n",
    "    print(f\"Memory: {r.value['memory']}\")\n",
    "    print(f\"Type: {r.value['type']}\\n\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
